import pandas as pd

# 封装训练分类器的函数，支持传入外部已训练好的分类器
import pandas as pd
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split

import matplotlib.pyplot as plt

# 设置支持中文的字体
plt.rcParams["font.sans-serif"] = ["Songti SC"]
plt.rcParams["axes.unicode_minus"] = False


def train_model(classifier):
    # 读取Excel文件
    df = pd.read_excel("data/processed_data_train.xlsx")
    # 划分特征和目标变量
    x = df.drop(["renewal", "policy_start_date", "policy_end_date"], axis=1)
    y = df["renewal"]
    # 标准化特征
    scaler = StandardScaler()
    x = scaler.fit_transform(x)
    # 划分训练集和测试集
    x_train, x_test, y_train, y_test = train_test_split(
        x, y, test_size=0.2, random_state=42
    )
    classifier.fit(x_train, y_train)
    # 模型评估
    y_pred = classifier.predict(x_test)
    accuracy = accuracy_score(y_test, y_pred)
    print(f"模型准确率: {accuracy}")
    print("分类报告:")
    print(classification_report(y_test, y_pred))

    # 如果是随机森林，绘制特征重要性柱状图
    if isinstance(classifier, RandomForestClassifier):
        feature_importances = classifier.feature_importances_
        feature_names = df.drop(
            ["renewal", "policy_start_date", "policy_end_date"], axis=1
        ).columns
        feature_importance_df = pd.DataFrame(
            {"特征": feature_names, "重要性": feature_importances}
        )
        # 按重要性降序排序
        feature_importance_df = feature_importance_df.sort_values(
            by="重要性", ascending=False
        )
        feature_importance_df.plot(x="特征", y="重要性", kind="bar", figsize=(10, 6))
        # 保存特征重要性图
        # 保存特征重要性图，文件名包含分类器名称
        plt.savefig(f"data/feature_importance_rf.png")
    # 如果是决策树，绘制决策树图
    elif isinstance(classifier, DecisionTreeClassifier):
        from sklearn.tree import plot_tree
        fig, ax = plt.subplots(figsize=(10, 10))  # 调整图片大小
        plot_tree(classifier, filled=True, feature_names=df.drop(
            ["renewal", "policy_start_date", "policy_end_date"], axis=1).columns, fontsize=10)  # 调整文字大小
        plt.tight_layout()  # 调整布局，避免文字重叠
        # 保存决策树图
        plt.savefig(f"data/decision_tree.png", dpi=300)  # 提高图片清晰度

    return classifier, scaler


from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.svm import SVC
from sklearn.naive_bayes import GaussianNB
from sklearn.metrics import accuracy_score, classification_report

classifiers = {
    # 随机森林
    "rf": RandomForestClassifier(n_estimators=100, random_state=42),
    # 逻辑回归
    "lr": LogisticRegression(random_state=42),
    # 决策树
    "dt": DecisionTreeClassifier(max_depth=3, random_state=42),
    # 支持向量机
    "svm": SVC(random_state=42),
    # 朴素贝叶斯
    "nb": GaussianNB(),
}

for name, classifier in classifiers.items():
    print(f"正在训练 {name} 分类器...")

    # 调用训练函数
    classifier, scaler = train_model(classifier)

    # 读取测试数据
    new_df = pd.read_excel("data/processed_data_test.xlsx")

    # 提取特征列
    new_x = new_df.drop(["policy_start_date", "policy_end_date"], axis=1)

    # 标准化特征
    new_x = scaler.transform(new_x)

    # 使用模型进行预测
    new_y_pred = classifier.predict(new_x)
    # 将new_y_pred的0和1转换为No和Yes
    new_y_pred = ["No" if pred == 0 else "Yes" for pred in new_y_pred]

    # 提取ID和预测结果
    result = pd.DataFrame({"ID": new_df["policy_id"], "renewal": new_y_pred})

    # 设置保存文件名：data/prediction_result_{name}.csv, name为分类器的name
    filename = f"data/prediction_result_{name}.csv"
    # 保存为CSV文件
    result.to_csv(filename, index=False)
